package utils

import (
	"bytes"
	"fmt"

	"github.com/klauspost/reedsolomon"
)

// Encode uses Reed-Solomon coding to split data into N shards,
// requiring K shards to reconstruct.
func Encode(k, n int, data []byte) ([][]byte, error) {
	enc, err := reedsolomon.New(k, n-k)
	if err != nil {
		return nil, fmt.Errorf("failed to create encoder: %w", err)
	}

	shards, err := enc.Split(data)
	if err != nil {
		return nil, fmt.Errorf("failed to split data: %w", err)
	}

	if err := enc.Encode(shards); err != nil {
		return nil, fmt.Errorf("failed to encode shards: %w", err)
	}

	return shards, nil
}

// Decode reconstructs the original data from a subset of shards.
// The shards slice can contain `nil` for missing shards.
func Decode(k, n int, shards [][]byte) ([]byte, error) {
	if len(shards) != n {
		return nil, fmt.Errorf("invalid number of shards provided: got %d, want %d", len(shards), n)
	}

	enc, err := reedsolomon.New(k, n-k)
	if err != nil {
		return nil, fmt.Errorf("failed to create decoder: %w", err)
	}

	// Try to reconstruct the data if verification fails.
	if ok, _ := enc.Verify(shards); !ok {
		err = enc.Reconstruct(shards)
		if err != nil {
			return nil, fmt.Errorf("reconstruction failed: %w", err)
		}

		// Verify again after reconstruction.
		ok, err = enc.Verify(shards)
		if err != nil || !ok {
			return nil, fmt.Errorf("verification failed after reconstruction")
		}
	}

	var buf bytes.Buffer
	// Calculate the original data length based on the first shard's length.
	if len(shards[0]) == 0 {
		return nil, fmt.Errorf("cannot determine data length from empty shards")
	}
	dataLen := len(shards[0]) * k
	if err := enc.Join(&buf, shards, dataLen); err != nil {
		return nil, fmt.Errorf("failed to join shards: %w", err)
	}

	return buf.Bytes(), nil
}
